import gym
import time
import matplotlib.pyplot as plt
from matplotlib import animation

env_name = "CartPole-v0"

def display():
    env = gym.make(env_name,render_mode="human")
    for episode in range(0,5):
        done = False
        score = 0
        env.reset()
        while not done:
            env.render()
            action = env.action_space.sample()
            observation,reward,done,info,_ = env.step(action)
            score+=reward
            time.sleep(0.2)
        print(f'episode: {episode},total reward: {score}')


def display_as_gif():
    frames = []
    score = 0
    #plt.rcParams['animation.writer'] = 'pillow'
    env = gym.make(env_name,render_mode = "rgb_array")
    env.reset()
    done = False
    while not done:
        frames.append(env.render())
        action = env.action_space.sample()
        obs,reward,done,info,_ = env.step(action)
        score+=reward
        time.sleep(0.1)    
    print(f'score: {score}')
    #渲染gif
    plt.figure(figsize=(frames[0].shape[1]/72,frames[0].shape[0]/72),dpi=72)
    patch = plt.imshow(frames[0])
    plt.axis('off')

    def animate(i):
        patch.set_data(frames[i])
    anim = animation.FuncAnimation(plt.gcf(),animate,frames=len(frames),interval=50)
    anim.save('E:/chartGPT_app/Reinforcement_learning/m_cartpole.gif') 
    env.close()


#display()
display_as_gif()
